#ifndef THREAD_POOL_H
#define THREAD_POOL_H

#include <condition_variable>
#include <thread>
#include <queue>
#include <atomic>
#include <functional>
#include <future>


class ThreadPool {
public:
    ThreadPool() : thread_num_(1), terminate_(false) {}
    virtual ~ThreadPool() { stop(); }

    bool init(size_t num)
    {
        if (!threads_.empty()) return false;

        thread_num_ = num;
        return true;
    }

    bool start()
    {
        std::lock_guard<std::mutex> lg(mutex_);

        if (!threads_.empty()) return false;
        for (size_t i = 0; i < thread_num_; ++i)
        {
            threads_.push_back( new(std::nothrow) std::thread(&ThreadPool::run, this));
        }

        return true;
    }

    void stop()
    {
        {
            std::lock_guard<std::mutex> lg(mutex_);
            terminate_ = true;

            condition_.notify_all();
        }

        for (auto&  td: threads_) {
            if (nullptr == td) continue;

            if (td->joinable()) td->join();

            delete td;
            td = nullptr;
        }

        std::lock_guard<std::mutex> lg(mutex_);
        threads_.clear();
    }

    void run()
    {
        while (!terminate_) {
            std::function<void()> task;
            bool ok = get(task);
            if (!ok) continue;

            ++atomic_;
            try {
                task();
            }catch(...) {
            }

            --atomic_;

            std::lock_guard<std::mutex> lg(mutex_);
            if (tasks_.empty() && atomic_ == 0) condition_.notify_all();
        }
    }

    template<typename T, typename... Args>
    auto exec(T&& t, Args&&... args) -> std::future<decltype(t(args...))>
    {
        using RetType = decltype(t(args...));

        auto task = std::make_shared<std::packaged_task<RetType()>>(std::bind(std::forward<T>(t), std::forward<Args>(args)...));
        std::future<RetType> result = task->get_future();

        std::lock_guard<std::mutex> lg(mutex_);
        tasks_.emplace([task] { (*task)(); });

        condition_.notify_one();
        return result;
    }

    bool waitForAllDone(int millsecond)
    {
        std::unique_lock<std::mutex> lg(mutex_);
        if (tasks_.empty()) return true;

        if (millsecond < 0)
        {
            condition_.wait(lg, [this] { return terminate_ || tasks_.empty(); });
            return true;
        }

        return condition_.wait_for(lg, std::chrono::milliseconds(millsecond), [this] { return terminate_ || tasks_.empty(); });
    }

protected:
    bool get(std::function<void()>& task)
    {
        std::unique_lock<std::mutex> lg(mutex_);
        if (tasks_.empty())
        {
            condition_.wait(lg, [this] { return !terminate_ || !this->tasks_.empty(); });
        }

        if (terminate_ || tasks_.empty()) return false;

        task = std::move(tasks_.front());
        tasks_.pop();

        return true;
    }

protected:
    size_t                              thread_num_;
    std::mutex                          mutex_;
    // 任务队列
    std::queue<std::function<void()>>   tasks_;
    // 工作线程
    std::vector<std::thread*>           threads_;
    std::condition_variable             condition_;
    volatile bool                       terminate_;
    std::atomic<int>                    atomic_{0};
};

#endif // THREAD_POOL_H
